import os
import time
from shutil import make_archive
import math
# pip install pycryptodome
from Crypto import Random
from Crypto.Cipher import AES
# pip install pillow
from PIL import  Image
import re 

def natural_sort_key(s):
    # 将字符串拆分为数字和非数字部分
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split(r'(\d+)', s)]

class CryGz:
    '''
    压缩，加解密文件
    示例：
        cry = CryGz()
        # cry.gen_keyiv() # 生成 key-iv
        # cry.encrypt_file('note.md', keyivfile='key-iv') # 加密 ‘note.md'
        cry.decrypt_file('note.md.cry', keyivfile='key-iv') # 解密 ’note.md.cry'
        # cry.gztar(org_dir) # 压缩目录
    '''

    def gztar(self, org_dir, fmt='zip'):
        ''' 
        将 org_dir 中的内容压缩
        
            注意：解压时需放入新建目录中，以防一盘散沙 
        '''
        name = os.path.basename(org_dir)
        basename = os.path.join(os.path.dirname(org_dir), name)
        make_archive(basename, format=fmt, root_dir=org_dir)
        print('生成压缩文件：', basename + '.' + fmt)

    def gen_keyiv(self, directory):
        now = time.strftime('%y%m%d-%H%M%S', time.localtime())
        key = Random.new().read(16)
        iv = Random.new().read(AES.block_size)
        with open(os.path.join(directory, 'key-iv'+now), 'wb') as f:
            f.write(key)
            f.write(iv)
        print(key,iv)

    def get_keyiv(self, filename):
        with open(filename, 'rb') as f:
            key = f.read(16)
            iv = f.read()
            return (key,iv)

    def aes_encyrpt(self, data, key, iv):
        aes = AES.new(key, AES.MODE_CFB, iv)
        crypt = aes.encrypt(data)
        return crypt

    def aes_decrypt(self, data, key, iv):
        aes = AES.new(key, AES.MODE_CFB, iv)
        decrypt = aes.decrypt(data)
        return decrypt

    def encrypt_file(self, filename, keyivfile='key-iv'):
        key, iv = self.get_keyiv(keyivfile)
        data = None
        with open(filename, 'rb') as f:
            data = f.read()
        cry_data = self.aes_encyrpt(data,key,iv)
        with open(filename+'.cry', 'wb') as wf:
            wf.write(cry_data)

    def decrypt_file(self, filename, keyivfile='key-iv'):
        key, iv = self.get_keyiv(keyivfile)
        with open(filename, 'rb') as f:
            data = f.read()
            decry_data = self.aes_decrypt(data, key, iv)
            with open(filename+'d', 'wb') as wf:
                wf.write(decry_data)

    def split_file(self, filename, count):
        if count <= 1: return
        size = os.path.getsize(filename=filename)
        size = math.ceil(size / count)
        with open(filename, 'rb') as f:
            i = 0
            while True:
                buffer = f.read(size)
                if not buffer: break
                i += 1
                root, ext = os.path.splitext(filename)
                tempname = f'temp-0{i}' if i < 10 else f'temp-{i}'
                outname = os.path.join(root, tempname + ext)
                dir = os.path.dirname(outname)
                if not os.path.exists(dir):
                    os.mkdir(dir)
                with open(outname, 'wb') as wf:
                    wf.write(buffer)

    def combine_files(self, dir, ext):
        outfile = os.path.join(dir, 'temp.out')
        paths = []
        for dirpath, dirnames, filenames in os.walk(dir):             
            for file in filenames:
                _, tempext = os.path.splitext(file)
                if ext and tempext != ext: continue
                paths.append(os.path.join(dirpath, file))
        paths = sorted(paths, key=natural_sort_key)

        with open(outfile, 'wb') as f:
            for path in paths:
                print(path)
                with open(path, 'rb') as rf:
                    buffer = rf.read()
                    f.write(buffer)

def blackwhite_image(filename, threshold=160):
    image = Image.open(filename)
    gray = image.convert("L")
    binary = gray.point(lambda x: 255 if x>threshold else x)
    dir = os.path.dirname(filename)
    name = os.path.basename(filename)
    temppath = os.path.join(dir,"黑白_"+name)
    binary.save(temppath)
    print("OK!")

def test():
    cry = CryGz()
    # cry.gen_keyiv()
    # cry.encrypt_file('note.md.tar.gz')
    # cry.decrypt_file('note.md.tar.gz.cry')
    # cry.gztar(os.path.dirname(__file__))
    # filename = os.path.join(os.path.dirname(__file__), 'temp/1.7z')
    # cry.split_file(filename, 3)
    # cry.combie_files(os.path.join(os.path.dirname(filename), '1'))

if __name__ == "__main__":
    test()
